import warnings

import mytorch
from utils import Timer
from dataset import Dataset
from config import workspace, trains, valids, train_epoch
from .run import Runner
from .log import Logger


def do_train(T: Timer = None):
    # 加载数据集 -> 选择目标数据集并指定模式
    if T: T.track(' -> loading trainset')
    # trainset = Dataset(names=trains)
    # !!!!!!!!!!!!!!!!!!!!!!!!
    trainset = Dataset(names=trains[:1])
    # train_loader = DataLoader(dataset=trainset, batch_size=batch_size, num_workers=load_worker)
    # validset = Dataset(names=valids)
    validset = Dataset(names=valids[:1])
    # valid_loader = DataLoader(dataset=validset, batch_size=batch_size, num_workers=load_worker)

    # 加载模型
    if T: T.track(' -> initializing net')
    model = mytorch.net.Jassor(pretrained=workspace / 'env' / 'van_base_828.pth.tar')

    # 训练项
    if T: T.track(' -> initializing train-listener')
    loss = mytorch.loss.CrossEntropy()
    optimizer = mytorch.optimizer.Adam(model)
    scheduler = mytorch.scheduler.Cosine(optimizer)

    # 现在 optimizer 和 scheduler 都将在 Epoch 中执行
    if T: T.track(' -> loading component-epoch')
    logger = Logger(
        model=model,
        metrics_seg=['loss', 'loss_seg', 'loss_aux', 'score_seg', 'dice', 'jaccard', 'liver', 'spleen', 'lkidney', 'rkidney'],
        metrics_cls=['loss', 'score_cls', 'recall_A', 'f1_A', 'acc_A', 'acc_B', 'f1_B', 'liver', 'spleen', 'lkidney', 'rkidney'],
        metrics_val=['score', 'score_cls', 'score_seg', 'recall_A', 'f1_A', 'acc_A', 'acc_B', 'f1_B', 'dice_D', 'jaccard_D',
                     'liver_C', 'spleen_C', 'lkidney_C', 'rkidney_C', 'liver_D', 'spleen_D', 'lkidney_D', 'rkidney_D'],
        T=T,
    )
    runner = Runner(
        model=model,
        loss=loss,
        optimizer=optimizer,
        scheduler=scheduler,
        logger=logger,
    )

    with logger:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            for epoch in range(train_epoch):
                logger.start_epoch(epoch)
                epoch += 1
                runner.train_seg(trainset, epoch)
                runner.train_cls(trainset, epoch)
                runner.valid(trainset, prefix='etr', epoch=epoch)
                runner.valid(validset, prefix='evd', epoch=epoch)
                logger.save_epoch(epoch)
